package com.xcf.weeklycontest198;

import sun.plugin.javascript.ocx.JSObject;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * SameNode <br>
 *
 * @author xiaocf <br>
 * @version 1.0 <br>
 * @date 2020/7/22 11:15 <br>
 */
public class SameNode {

    /**
     * 给你一棵树（即，一个连通的无环无向图），这棵树由编号从 0  到 n - 1 的 n 个节点组成，且恰好有 n - 1 条 edges 。树的根节点为节点 0 ，树上的每一个节点都有一个标签，也就是字符串 labels 中的一个小写字符（编号为 i 的 节点的标签就是 labels[i] ）
     * <p>
     * 边数组 edges 以 edges[i] = [ai, bi] 的形式给出，该格式表示节点 ai 和 bi 之间存在一条边。
     * <p>
     * 返回一个大小为 n 的数组，其中 ans[i] 表示第 i 个节点的子树中与节点 i 标签相同的节点数。
     * <p>
     * 树 T 中的子树是由 T 中的某个节点及其所有后代节点组成的树。
     * <p>
     * 例1：
     * 输入：n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], labels = "abaedcd"
     * 输出：[2,1,1,1,1,1,1]
     * 解释：节点 0 的标签为 'a' ，以 'a' 为根节点的子树中，节点 2 的标签也是 'a' ，因此答案为 2 。注意树中的每个节点都是这棵子树的一部分。
     * 节点 1 的标签为 'b' ，节点 1 的子树包含节点 1、4 和 5，但是节点 4、5 的标签与节点 1 不同，故而答案为 1（即，该节点本身）。
     * <p>
     * 例2：
     * 输入：n = 4, edges = [[0,1],[1,2],[0,3]], labels = "bbbb"
     * 输出：[4,2,1,1]
     * 解释：节点 2 的子树中只有节点 2 ，所以答案为 1 。
     * 节点 3 的子树中只有节点 3 ，所以答案为 1 。
     * 节点 1 的子树中包含节点 1 和 2 ，标签都是 'b' ，因此答案为 2 。
     * 节点 0 的子树中包含节点 0、1、2 和 3，标签都是 'b'，因此答案为 4 。
     *
     *
     * <p>
     * 来源：力扣（LeetCode）
     * 链接：https://leetcode-cn.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label
     * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。
     *
     * @param n      树的节点数
     * @param edges  边数组
     * @param labels 节点标签列表
     * @return 相同子节点数列表
     * <p>
     * 暴力解题，时间复杂度O(n^2) 时间超时
     */
    public int[] countSubTrees(int n, int[][] edges, String labels) {
        if (n <= 1) {
            return new int[]{1};
        }
        int[] result = new int[n];
        Map<Integer, Integer> parentMap = new HashMap<>( );
        result[0] = 1;
        for (int[] edge : edges) {
            int parent = edge[0];
            int currNode = edge[1];
            if (currNode == 0 || (parent != 0 && (parentMap.containsKey(currNode) && !parentMap.get(currNode).equals(parent)))) {
                // 父节点
                parent = edge[1];
                currNode = edge[0];
            }
            result[currNode] = 1;
            parentMap.put(currNode, parent);
            char charAt = labels.charAt(currNode);
            while (parentMap.containsKey(currNode)) {
                currNode = parentMap.get(currNode);
                if (labels.charAt(currNode) == charAt) {
                    result[currNode] = result[currNode] + 1;
                }
            }
        }
        return result;
    }

    public int[] countSubTreeV2(int n, int[][] edges, String labels) {
        if (n <= 1) {
            return new int[]{1};
        }
        Map<Integer, List<Integer>> tree = new HashMap<>( );
        for (int[] edge : edges) {
            int node0 = edge[0];
            int node1 = edge[1];
            List<Integer> list0 = tree.getOrDefault(node0, new ArrayList<>( ));
            list0.add(node1);
            tree.put(node0, list0);
            List<Integer> list1 = tree.getOrDefault(node1, new ArrayList<>( ));
            list1.add(node0);
            tree.put(node1, list1);
        }
        int[] result = new int[n];
        boolean[] visited = new boolean[n];
        depthSearch(0, visited, tree, result, labels);
        return result;
    }

    private int[] depthSearch(int currNode, boolean[] visited, Map<Integer, List<Integer>> tree,
                             int[] result, String labels) {
        visited[currNode] = true;
        int[] currCount = new int[26];
        char currChar = labels.charAt(currNode);
        currCount[currChar - 'a']++;
        List<Integer> child = tree.get(currNode);
        for (Integer next : child) {
            if (!visited[next]) {
                int[] childCount = depthSearch(next, visited, tree, result, labels);
                for (int i = 0; i < childCount.length; i++) {
                    currCount[i] = currCount[i] + childCount[i];
                }
            }
        }
        result[currNode] = currCount[currChar - 'a'];
        return currCount;
    }

    /**
     * 官方解法
     */
    public int[] countSubTreesDepth(int n, int[][] edges, String labels) {
        Map<Integer, List<Integer>> edgesMap = new HashMap<>( );
        for (int[] edge : edges) {
            int node0 = edge[0], node1 = edge[1];
            List<Integer> list0 = edgesMap.getOrDefault(node0, new ArrayList<>( ));
            List<Integer> list1 = edgesMap.getOrDefault(node1, new ArrayList<>( ));
            list0.add(node1);
            list1.add(node0);
            edgesMap.put(node0, list0);
            edgesMap.put(node1, list1);
        }
        int[] result = new int[n];
        boolean[] visited = new boolean[n];
        depthFirstSearch(0, result, visited, edgesMap, labels);
        return result;
    }

    public int[] depthFirstSearch(int node, int[] result, boolean[] visited, Map<Integer, List<Integer>> edgesMap, String labels) {
        visited[node] = true;
        // 当前节点及子节点下26个字母的统计情况
        int[] curCounts = new int[26];
        curCounts[labels.charAt(node) - 'a']++;
        List<Integer> nodesList = edgesMap.get(node);
        for (int nextNode : nodesList) {
            if (!visited[nextNode]) {
                // 搜索子节点26个字母的统计情况
                int[] childCounts = depthFirstSearch(nextNode, result, visited, edgesMap, labels);
                for (int i = 0; i < 26; i++) {
                    // 合并子节点的情况
                    curCounts[i] += childCounts[i];
                }
            }
        }
        // 当前节点的字母统计情况：当前节点+子节点字母的统计情况
        result[node] = curCounts[labels.charAt(node) - 'a'];
        // 并且将结果返回给父节点
        return curCounts;
    }


    /**
     *    1 4 5
     * 0
     *    2  3 6
     * @param args
     */
    public static void main(String[] args) {
        SameNode sameNode = new SameNode( );
        int[] trees = sameNode.countSubTreeV2(7, new int[][]{{0, 1}, {0, 2}, {1, 4}, {1, 5}, {2, 3}, {2, 6}}, "abaedcd");
        for (int tree : trees) {
            System.out.print(tree + ", ");
        }
        System.out.println( );
    }

    private static void excute(int n, String arra, String labels) {
        SameNode sameNode = new SameNode( );
        int[] trees = sameNode.countSubTreesDepth(n, new int[][]{{0, 1}, {0, 2}, {1, 4}, {1, 5}, {2, 3}, {2, 6}}, "abaedcd");
        for (int tree : trees) {
            System.out.print(tree + ", ");
        }
        System.out.println( );
    }
}
